{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "https://langchain-ai.github.io/langgraph/how-tos/pass-run-time-values-to-tools/\n",
    "Sometimes, you want to let a tool-calling LLM populate a *subset* of the tool functions' arguments and provide the other values for the other arguments at runtime. If you're using LangChain-style [tools](https://python.langchain.com/docs/concepts/#tools), an easy way to handle this is by annotating function parameters with [InjectedArg](https://python.langchain.com/docs/how_to/tool_runtime/). This annotation excludes that parameter from being shown to the LLM.\n",
    "\n",
    "In LangGraph applications you might want to pass the graph state or [shared memory](https://langchain-ai.github.io/langgraph/how-tos/cross-thread-persistence/) (store) to the tools at runtime. This type of stateful tools is useful when a tool's output is affected by past agent steps (e.g. if you're using a sub-agent as a tool, and want to pass the message history in to the sub-agent), or when a tool's input needs to be validated given context from past agent steps.\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7610f37cbbf8f474"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# from typing import Annotated, List\n",
    "# \n",
    "# from langchain_core.runnables import RunnableConfig\n",
    "# from langgraph.prebuilt import InjectedState, InjectedStore\n",
    "# from langgraph.prebuilt.chat_agent_executor import AgentState\n",
    "# from langgraph.store.base import BaseStore\n",
    "# \n",
    "# # Can be sync or async; @tool decorator not required\n",
    "# async def my_tool(\n",
    "#         # These arguments are populated by the LLM\n",
    "#         some_arg: str,\n",
    "#         another_arg: float,\n",
    "#         # The config: RunnableConfig is always available in LangChain calls\n",
    "#         # This is not exposed to the LLM\n",
    "#         config: RunnableConfig,\n",
    "#         # The following three are specific to the prebuilt ToolNode\n",
    "#         # (and `create_react_agent` by extension). If you are invoking the\n",
    "#         # tool on its own (in your own node), then you would need to provide these yourself.\n",
    "#         store: Annotated[BaseStore, InjectedStore],\n",
    "#         # This passes in the full state.\n",
    "#         state: Annotated[State, InjectedState],\n",
    "#         # You can also inject single fields from your state if you\n",
    "#         messages: Annotated[list, InjectedState(\"messages\")]\n",
    "#         # The following is not compatible with create_react_agent or ToolNode\n",
    "#         # You can also exclude other arguments from being shown to the model.\n",
    "#         # These must be provided manually and are useful if you call the tools/functions in your own node\n",
    "#         # some_other_arg=Annotated[\"MyPrivateClass\", InjectedToolArg],\n",
    "# ):\n",
    "#     \"\"\"Call my_tool to have an impact on the real world.\n",
    "# \n",
    "#     Args:\n",
    "#         some_arg: a very important argument\n",
    "#         another_arg: another argument the LLM will provide\n",
    "#     \"\"\"  # The docstring becomes the description for your tool and is passed to the model\n",
    "#     print(some_arg, another_arg, config, store, state, messages)\n",
    "#     # Config, some_other_rag, store, and state  are all \"hidden\" from\n",
    "#     # LangChain models when passed to bind_tools or with_structured_output\n",
    "#     return \"... some response\"\n",
    "# \n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:52.580600Z",
     "start_time": "2024-11-11T03:50:52.574703Z"
    }
   },
   "id": "96e887068229798e",
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Pass graph state to tools"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "13045b7579c7f5e3"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# this is the state schema used by the prebuilt create_react_agent we'll be using below\n",
    "from langgraph.prebuilt.chat_agent_executor import AgentState\n",
    "from langchain_core.documents import Document\n",
    "\n",
    "# Agent state schema增加一个入参docs\n",
    "class State(AgentState):\n",
    "    docs: List[str]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:54.113765Z",
     "start_time": "2024-11-11T03:50:52.609954Z"
    }
   },
   "id": "36f33436fabaab26",
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Define the tools\n",
    "We'll want our tool to take graph state as an input, but we don't want the model to try to generate this input when calling the tool. We can use the InjectedState annotation to mark arguments as required graph state (or some field of graph state. These arguments will not be generated by the model. When using ToolNode, graph state will automatically be passed in to the relevant tools and arguments.\n",
    "\n",
    "In this example we'll create a tool that returns Documents and then another tool that actually cites the Documents that justify a claim."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "73b541cfc5ee0583"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "\n",
    "from typing_extensions import Annotated\n",
    "\n",
    "from langchain_core.tools import tool\n",
    "from langgraph.prebuilt import InjectedState\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_context(question: str, state: Annotated[dict, InjectedState]):\n",
    "    \"\"\"Get relevant context for answering the question.\"\"\"\n",
    "    return \"\\n\\n\".join(doc for doc in state[\"docs\"])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:54.131041Z",
     "start_time": "2024-11-11T03:50:54.117498Z"
    }
   },
   "id": "e4e3ba6a45c8f990",
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "如果我们查看这些工具的输入模式，我们会发现状态仍然被列出："
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "417d1063ceb61134"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "{'description': 'Get relevant context for answering the question.',\n 'properties': {'question': {'title': 'Question', 'type': 'string'},\n  'state': {'title': 'State', 'type': 'object'}},\n 'required': ['question', 'state'],\n 'title': 'get_context',\n 'type': 'object'}"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_context.get_input_schema().schema()\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:54.145942Z",
     "start_time": "2024-11-11T03:50:54.133957Z"
    }
   },
   "id": "2f5d4fda6238fc0d",
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "但是如果我们查看工具调用模式，这是传递给模型进行工具调用的，状态已被删除："
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "274bf03b13cf4876"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "{'description': 'Get relevant context for answering the question.',\n 'properties': {'question': {'title': 'Question', 'type': 'string'}},\n 'required': ['question'],\n 'title': 'get_context',\n 'type': 'object'}"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_context.tool_call_schema.schema()\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:54.157095Z",
     "start_time": "2024-11-11T03:50:54.150449Z"
    }
   },
   "id": "63f2300eec5f6f74",
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Define the graph\n",
    "In this example we will be using a [prebuilt ReAct agent](https://langchain-ai.github.io/langgraph/how-tos/create-react-agent/). We'll first need to define our model and a tool-calling node ([ToolNode](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.ToolNode)):"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8da35f6a00e41391"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/4gHYSUNDX1BST0ZJTEUAAQEAAAHIAAAAAAQwAABtbnRyUkdCIFhZWiAH4AABAAEAAAAAAABhY3NwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAA9tYAAQAAAADTLQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAlkZXNjAAAA8AAAACRyWFlaAAABFAAAABRnWFlaAAABKAAAABRiWFlaAAABPAAAABR3dHB0AAABUAAAABRyVFJDAAABZAAAAChnVFJDAAABZAAAAChiVFJDAAABZAAAAChjcHJ0AAABjAAAADxtbHVjAAAAAAAAAAEAAAAMZW5VUwAAAAgAAAAcAHMAUgBHAEJYWVogAAAAAAAAb6IAADj1AAADkFhZWiAAAAAAAABimQAAt4UAABjaWFlaIAAAAAAAACSgAAAPhAAAts9YWVogAAAAAAAA9tYAAQAAAADTLXBhcmEAAAAAAAQAAAACZmYAAPKnAAANWQAAE9AAAApbAAAAAAAAAABtbHVjAAAAAAAAAAEAAAAMZW5VUwAAACAAAAAcAEcAbwBvAGcAbABlACAASQBuAGMALgAgADIAMAAxADb/2wBDAAMCAgMCAgMDAwMEAwMEBQgFBQQEBQoHBwYIDAoMDAsKCwsNDhIQDQ4RDgsLEBYQERMUFRUVDA8XGBYUGBIUFRT/2wBDAQMEBAUEBQkFBQkUDQsNFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBT/wAARCAD5ANYDASIAAhEBAxEB/8QAHQABAAICAwEBAAAAAAAAAAAAAAYHAwUCBAgBCf/EAFIQAAEEAQIDAgUOCQkGBwAAAAEAAgMEBQYRBxIhEzEWFyJBlAgUFTJRVVZhcXSy0dLTIzY3QlSBkZOVGDVDUnWCkrO0JCUncpahMzRTZLHB8P/EABsBAQEAAwEBAQAAAAAAAAAAAAABAgMFBAYH/8QAMxEBAAECAQkFCAIDAAAAAAAAAAECEQMEEiExQVFSkdEUM2FxoQUTFSNiscHhgZIi8PH/2gAMAwEAAhEDEQA/AP1TREQEREBERAWG1cr0o+exPHXZ/WleGj9pWju37uevz47FTGlVrnkt5NrQ5zX/APpQhwLS4d7nuBa3cNAc4u5Ptbh/p+F5llxcF+ydua1fb65mcR5y9+5/Z0W+KKae8n+IW293fCrC++9D0ln1p4VYX34oeks+tPBXC+89D0Zn1J4K4X3noejM+pX5Pj6LoPCrC+/FD0ln1p4VYX34oeks+tPBXC+89D0Zn1J4K4X3noejM+pPk+PoaDwqwvvxQ9JZ9aeFWF9+KHpLPrTwVwvvPQ9GZ9SeCuF956HozPqT5Pj6Gg8KsL78UPSWfWu5UyFW+0uq2YbLR3mGQOA/Yun4K4X3noejM+pdS1oHTluQSuw1OGdp3bYrRCGZp+KRmzh+op8mds+n6TQ36KMR2bmkZ4Yb9qbJYeVwjZen5e1quJ2a2UgAOYegD9twdubfcuEnWuujN8YJgREWtBERAREQEREBERAREQEREBajV2Yfp/S+VyMQDpq1Z8kTXdxft5IP69lt1HuIVOW9onMxwtMkza7pWMaNy5zPLAA90luy24MROJTFWq8LGtsNP4ePAYapQjPN2LPLk88khO73n43OLnE+6StisNO1FeqQWYHc8MzGyMd7rSNwf2FZlhVMzVM1a0FEuIHFbS3C6LHv1JkzSfkJHRVIIa01madzW8z+SKFj3kNHUnbYbjchS1Up6pWhUfBp3Jx4/WDdSY59mTEZzR2ON2ahK6NocyaIBwdHL0Ba5paeXqW9CsR2cp6pjT+N4q6b0m2tetUc3hfZeHJ1cdbnB55IWwtDY4XeS5sjnOkJAZs0O5S4KQWuP2gqOuW6Qs571vnX2m0WxS052wmw4bthE5j7LtDuNm8+53A2VUx5fWendd8Ltfax0nlrtuxpGzicxDp6g+4+neklrTDnij3LWu7J43G4aehPnUA4t4/Wep5tTDMYbX+W1Bj9VwW8fUxsEwwsOJguRSRyRtjIjsSGJpJGz5ec9GgDoHpi3x20TT1je0ocpYsahozR17VCnjbVh8DpI2yMLzHE4NYWvb5ZPLuSN9wQNXwF4943jngrNyrRu465XsWY5K89KyyMRssSRRubNJExj3OawOcxpJYSWuAIXW4S6fu4zjFxpyVrG2KkGSy2PdVtzQOY21GzHQNJY4jZ7Wv529NwDzDv3Wr9THYyGl8PlNCZjT2axuSxeUylr19YovbQswy3pJY3Q2NuR5c2Zp5Qdxyu3A2QXgiIg6+QoV8rQs0rcTZ6tmN0MsT+57HDZwPyglajQ1+e/puEWpe3t1JZqM0p33kfDK6IvO/9bk5v1rfqM8PG9pp+S4N+S/dtXI+YbbxyTvdGdvjZyn9a9FPc1X3x+V2JMiIvOgiIgIiICIiAiIgIiICIiAiIgilOdmg3mjb2iwDnl1O315Km53MMp7mN3J5H9G7bMOxDe0x6r4RaG1/kY8lqPSWEz95sQhZayFGKeQRgkhoc4E8u7nHb4ypa9jZGOY9oexw2LXDcEe4VGn8PsdCScbZyGFB/osdbfHEPc2iO8bf1NH/YL0TVRiaa5tPO/wDv8stEo8fU28KC0N8W+luUEkD2Jg2B8/5vxBSbR/DvS3D2GzFpjT2M0/FZc107MbUZAJSNwC4NA323Pf7qw+BNj4VZ799D90ngTY+FWe/fQ/dJ7vD4/SUtG9KEUX8CbHwqz376H7pRO9jstX4q4PTzNU5j2OuYW/flJlh7TtYZ6bGbfg/a8tiTfp38vUed7vD4/SS0b1qLS6s0XgNd4xuO1HhaGdx7ZBM2rka7Z4w8AgO5XAjcBxG/xldHwJsfCrPfvofuk8CbHwqz376H7pPd4fH6SWje0DfU3cKWBwbw40u0PGzgMTB1G4Ox8n3QP2LZ6Z4K6A0Zl4srgNF4HDZOIObHco4+KGVocNnAOa0EbgkFdzwJsfCrPfvoful98AKdh3+8MhlcqzffsbV14iPysZytcPicCEzMONdfKP8AhaHHK5Dwu7fDYqXnqP5ochkYXeRCzqHRRuHfKe7p7QbuJB5WuksEEdaCOGFjYoo2hjGMGwa0DYADzBfKtWGlXjr14Y68EbQ1kUTQ1rQO4ADoAsqwrriYzadUEiIi1IIiICIiAiIgIiICIiAiIgIiICIiAiIgKvssW+P7SwJPN4MZfYebb11jd/P8nm/WPPYKr/K7+P7S3Vu3gxl+hA3/APNY3u8+3ydO7fzILAREQEREBERAREQEREBERAREQEREBERAREQEREBERAREQFXuWA/lA6VPM0HwXzHk7dT/ALXjOu+3d+vzj9VhKvctt/KC0r1PN4L5jYcv/u8Z5/8A9/2QWEiIgIiICIiAiIgIiICIiAiIgIiICIiAiIgIiICIonf1ZkbVyxBg6NazFXkMMtu7O6JhkG4c1gaxxdykbE9ADuBuQdtuHh1Yk2pW10sRQj2d1h+gYP0ub7tPZ3WH6Bg/S5vu1v7LXvjnBZN14D1j6vbK6e9URXxNrhXO7UOJjuadGPizAd28s9is5r2O9b78p9bjbYeUHg+YL2L7O6w/QMH6XN92qgz3qf5tQ+qDw/Fqxj8MMzjqvYmoLEhinmaOWKdx7PfnY07D/lZ/V6uy1745wWelkUI9ndYfoGD9Lm+7T2d1h+gYP0ub7tOy1745wWTdFCPZ3WH6Bg/S5vu1li1flsW5kmdoU4qBcGvtUbD5OwJOwc9jmDyN9t3AnbfcjYFwk5LibLT/ADBZMkRF5EEREBERAREQEREBERAREQEREBERAVeaGO+BeT3m/eJ+M+upVYarzQv8wP8An13/AFUq9+T93V5x+V2JAiItiCIiAiLo2M5j6uXqYua7BHkrcckteo6QCWVjOXnc1veQ3mbufNzD3UHeUd4jnbh7qg9Nxi7RG43/AKJykSjnEj8neqf7Ktf5LluwO9o84+7KnXCxGe1HyLkuLPaN+RclxmIiIgIiICIiAiIgIiICIiAiIgIiICrzQv8AMD/n13/VSqw1Xmhf5gf8+u/6qVe/J+7q84/K7EgXkPiHrLUMOqb+t9KXNSMw2K1ZVw1qfI6gIpTO9dx1rEEOOEZa6Pdzm9o5zXhwLhuAvXirXOepw4dajyOSvZDTgnnyMxtWGtuWGRmc7bzsjbIGRzdP/FYGv7/K6lWqJnUig+Kua1DndU8QMQNRarp68hzFSrpvT2IsWIaU+NeIfwjhFs0hwNkvlc4FnJ0LdgDtci/iXxc1xxGOCuWKR0/lX4jHMg1XLi2U+SGNzJpKrKsrbAe55fvI4gjyQG8u5kHE/wBTzq3VeuM7k9PS4fADJyxyx52tm8tWu1ntjYwymrFIK80gDBsTyggNDgdtzZ2p+AGhda5o5jOYQXctJCyC1aiszV/XjWDZonZE9rZQPceHbDp3LDNmbiqW4jU+tNea+xeoNW5vGXcLpfEWew0/k5a1aO/JDZ7WVnLsS3ni6NOzXD2zSQNtHgqLuK3ELgJn83lMvDksroqzasyY7KT0w+ZgqOJAie0DmMji4Do4BoO4a3b0xDofCQZzN5iOly5HNVoad+btX/hoog8Rt5ebZuwlf1aATzdSdhtHstwH0Nm9OadwdvCE4/T0XY4oQ3LEU1WPkDC1szJBIQWgAguPNsN99llmyJ8o5xI/J3qn+yrX+S5SMDYAKOcSPyd6p/sq1/kuXqwO9o84+7KnXCxGe0b8i5Liz2jfkXJcZiIiICIiAiIgIiICIiAiIgIiICIiAq80L/MD/n13/VSqw1XczMhpjMzY7HYufO0rEs9thqPa19RznCSSKQyFrBu6YFg5g4tcQG7Rlx92TzGbVRe0zadOjVfqsarJCi0nstnvgZlfSqX36ey2e+BmV9Kpffr05n1R/aOq2btFpPZbPfAzK+lUvv1F7vGOtj+IWP0PYwd+LVWQqPu1scZ6vNJCzfmdzdtyjucdidyGkgbApmfVH9o6llhotJ7LZ74GZX0ql9+nstnvgZlfSqX36Zn1R/aOpZu1HOJH5O9U/wBlWv8AJcux7LZ74GZX0ql9+sWQx+e1Tj56EmElxVWaMtsOtWYjJIzY7xs7NzgHO9rzEgNDidiRsc8O2HXFdVUWib646kRabrBZ7RvyLktZhs/Xy7WRFrqWSFeKxYxdl7PXNVsnNyiRrHOA6se3mBLSWO5XHZbNcViIiICIiAiIgIiICIiAiIgIiICL45wY0ucQ1oG5J7gtDG+xqew2SOSaliIJz7URublIzF0IduS2Lmee7lc50QIPZn8IHGfIWdSiatiZZadMxwyszkXZSRSgyeXHCNyS7kad3lvKO0YW85Dg3bY3FU8PDJDRqxVIpJpLD2xMDQ6SR5fI87d7nOcST5ySs1atDSrRV68TIIImCOOKJoa1jQNg0AdAAOmyyoCIiAvzx4g+pl43Z71XVTWVbUWlaufnM2ZxcbrtoxQVKksEQgeRX84sRggAg7v3Pu/ocq/yHLNx8wHKGl1fTOR5zueZoktUeXp3bHsnf4flQWAiIgIiINbmcFBmIXDtZqVrZoZepuDJ4w17XgB2x8kuY3dpBa4dHAgkLpw5y5jrorZuGGIWrskNCxSEkkb4gznZ2/k7Qv6Pb1cWuLAQ4OkEY3y+OaHtLXAOaRsQe4oPqKMCrNoam31jBLa05SqNiZjasTprUJEnVzCXbvYI3H8GAXARAMDiQ1SSOVkrS5j2vaCW7tO43B2I/UQR+pBzREQEREBERAREQEREBEWK1P61rTTcj5ezYX8kY3c7Yb7AecoNBZEOsr1zHu5J8JUdJTyVK5j+eO690bHBjXv8l0bQ883K1wL9m8wMcjDJFodBx8mi8I7tcpMZKkcxfmz/ALbu9ocRMB0DxzbFo6AjYdAFvkBERAREQFX3DgnVeodQa435qOREWOxDt9w+jAXkTjrttLLLM4Ee2jbCfc256ltS8QsrY0pjJnR4iu8Mz+Qhc5ruXYO9ZROHdI8Edo4Hdkbths+RrmTqvXiqQRwQRshhiaGMjjaGtY0DYAAdwA8yDIiIgIiICIiAo9fqeC5tZShEGUS+W7kadeo+eaw7kA54g078/kAloa4v67DmO5kKIMdexHbrxTwvEkUrQ9jx3OaRuCsi0OBgmxeZy2O7C++kXNvQ3bdgTRudM+TtII9zzNDCwO5T0AlaGnYcrd8gIiICIiAiIgIi0uY1tp7T9oVsnnMdj7JHN2Nm0xj9vd5Sd9lnTRVXNqYvK2u3SKLeNLR3wpxHpsf1qM8S7/DbivoTM6Sz+o8VNispB2MoZfja9pBDmPad/bNe1rhv03aNwR0W3s+NwTylc2dzY6F4gaXhlqaMOpN9TUnS0his7kInZicQlw7Z8fNzvD42CVr9vKjc157yp8vzi9RTwXo8FfVE6vv6jzeLkx+Hpmticp65YIrhmcPwkZ323EbXBw72l+x+P3p40tHfCnEemx/WnZ8bgnlJmzuSlFFvGlo74U4j02P608aWjvhTiPTY/rTs+NwTykzZ3JSobns7kNQZeTTmm5ewkiLRlczy8zcewjfsotxyvsub3NO4ia4SPB3jjm1GS4jVdZ51ml9LZypA+WPnt5eKeNzoWEe0rNduJZj7uxZGOrtzysdOsHg6Gm8XDjsbWbVpw8xbG0kkuc4ue9zjuXOc5znOc4lznOJJJJK1VUVUTauLJaz5gcDQ0xiK2MxlcVqVcEMZzFxJJLnOc5xLnvc4lznuJc5ziSSSStgiLBBERAREQEREBERBHrVH/iDjbjcZPJ/uu1E/JNsbRQ/ha5bC6L85z/KcHfmiJw/OUhVMZT1QHCqHibh3y690xzwYvIQvu+EtVsNdxmp7wyR9p1kfyktcerRDIPzlc6AiIgIiICIiDpZq47H4e9aYAXwQSStB91rSR/8ACiOkqkdbAUpAOaezEyeeZ3V80jmgue4nqSSf1d3cFJ9VfixmPmc30Co9pr8XMV80i+gF0MDRhT5rsbJERZoIiICIiDq5LG1stTkrWoxJE/49i0jqHNI6tcDsQ4dQQCOq7+g8pPmtF4O9af2tmenE+WTbbndyjd23m3PXb41iWHhZ+TnTnzGL6KxxdODPhMfaei7EpREXOQREQERRvXWs4NFYgWHRizcnf2VWrzcvav7ySfM1o3JPuDYbkgHZh4dWLXFFEXmRucnlqOEqOt5G5XoVW+2ntStjYPlc4gKMS8YdHQvLTnIXEdN445Hj9oaQqPydq1ncj7IZWw6/e68skg8mIb+1jb3Mb0HQdTsCST1WNfW4XsPDin5tc38P3cvC8fHNo336b6PL9hPHNo336b6PL9hUci3fA8m4qucdC8KC4kep00nqn1Y2O1JXuRnh7kpPZjKuEUgbHYYd3wcu3N+FfynoNgHu9xe7vHNo336b6PL9hUcifA8m4qucdC8Lx8c2jffpvo8v2F9Zxk0a923s3G343wyNH7S1UaifA8m4qucdC8PS2H1BjNQ13T4vIVchE08rnVpWyBp9w7HofiK2C8sQGSlejvUp5KN+P2lquQ17fiPQhw6DyXAg7dQVevDfXw1jSmr22sgy9MNE8bPaytPdKweZpIII72kEdRsTxcu9l1ZLT7yib0+sLr1JkiIuEjV6q/FjMfM5voFR7TX4uYr5pF9AKQ6q/FjMfM5voFR7TX4uYr5pF9ALo4Pcz5/hdjvWHSMgkdCxsswaSxjncoc7boCdjt18+xXnbhbx61RjOCuY1nrzFRWK9S9bgqzY+6JrN2f2Qkrx1hD2MbWbO5I2u5jzAcxDeq9Grz3DwC1dLoHUugp8jhYsA6/Nl8DloTK65DZN4XImzxFoZyteXNJa8kjboFJvsRIG+qEn0tazNTiHpg6QtUMLLn4vWuQbkI7NaJwbK1rwxm0rXOYOTbY842cQsFfjfnZ7FXEan0dNo6bUGLt2sJZjybbTnvih7V0UoaxphlDDzgAuHku8rcLW5ngRqji5kM3e4i3MNRdPp2xp+hU086WaOHt3NdJZe+VrCXbxx7MA2AB3J713cdwo11q/VWmsjr+/gmVNNU7UNRmBMz33LE8Brunl7RrRGBGX7MbzdXnyugU/yGj0lxxzGmuGHBbGRYt2q9UarwjJmz5XLCoyR8UETpOad7Xl8rzINm7Eu2cSRsvQmPmns0K01msadmSJr5a5eH9k8gEs5h0Ox3G46HZefrHBbXzuCGB4e2KOhdRV8fUkx0kmV9ctHZsa1lWxHyscWTNAcXAefbleFdmg9P29KaJwGFv5KTMXsdQgqT5CbfnsvZGGukO5J3cQT1JPXqSrTfaN6sPCz8nOnPmMX0VmWHhZ+TnTnzGL6KuL3M+cfaV2JSiIucgiIgKguLOSdkuIliBziYsbVjgjae5rpPwjyPlHZA/8gV+qguLONdjOIc87mkRZOrHPG89znx/g3gfIOyP98Lvexc3tWnXaben4uuyUWRdfI34sXRntziUwwsL3iGF8r9h7jGAucfiAJUVHFvT5/os5/wBO5D7hfb1YlFGiqYhrTJzg1pJIAHUk+ZUnS9VBh7uQqPZBjzhLdtlSKdmagde8p/I2R1MeWGFxB9sXBp3LQp2zijp++9tXsc0e3PZ7P0/fY079OrjAAB17ydlHuH2hNXaDix+n2v0/e0zQkc2K9M2UX3V9yWsLAOTmG4HPzdw9ruvJiV111U+5q0bbWndb8qxT8br9eHKZKTSxbp7F5mTD3L/sg3tGltgQiVkXJ5Td3NJBc0jcgcwG56/EzihmJsPrmjpfCTXIMLRniu5pt8VjVnMBftCNiXvja5rjsW7HoDus+R4TZe3w61hgGWaQuZjOzZOu9z39m2J9tkwDzybh3K0jYAjfz+dYNQ8NNYV/DnH6cs4WTCaqE00gybpmTVbEsAikLeRpD2u5Wnrtsfd8+iqcozbTfTHhfb+hY+i55bWjsFNNI+aaShA98kji5znGNpJJPeSfOtwoLj9b4rRuMoYO+3KSXcfWhrTOp4W9PEXNjaCWyMhLXD4wVn8bunj/AEWd/wCnch9wvbTi4cRETVF/NEzW20VknYfXuAsscWiac0pQPz2StIA/xiN391RvC5qtn8dHdqCw2B5IAtVpa8nQ7HdkjWuHd5x1Uk0TjXZnXuArMbzNgnN2Uj8xkbSQf8ZjH95TKJonArmrVafsyp1vSCIi/MFavVX4sZj5nN9AqPaa/FzFfNIvoBSnM03ZHEXqjCA+eCSIE+YuaR/9qIaSuR2MDThB5LNaFkFiB3R8MjWgOY4HqCD+0bEdCF0MDThTHiuxuERFmgiIgIiICw8LPyc6c+YxfRWPJ5StiKj7NqURxt6Ad7nuPQNa0dXOJIAaNySQB1K2GhMXPhNGYSjaZ2dmCnEyWPffkfyjdu/n2PTf4lji6MGfGY+09V2N6iIucgiIgKOa50ZBrXDis+QVrcL+1q2uXmMT+7qOm7SNwRv3HoQQCJGi2YeJVhVxXRNpgeXcrUtafyHrDLVzj7nXla87slH9aN/c8d3d1G43DT0WNenMli6WZqPq36kF6s/20NmJsjD8rSCFGJeEGjpXFxwNdpPXaNz2D9gIC+twvbmHNPzaJv4fstCikV5eJvRvvHF+9k+0nib0b7xxfvZPtLd8cybhq5R1LQo1FeXib0b7xxfvZPtJ4m9G+8cX72T7SfHMm4auUdS0KNRXl4m9G+8cX72T7S+s4O6NY7f2Cgd8T3vcP2F2yfHMm4auUdS0b1F1hLkLzKNGCS/ff7WrXAc8/GeuzR1HlOIA36lXtw40ENG0Zp7T2T5e3ymeRntI2j2sTD3loJJ3PVxJOwGzWyLEYLG4CuYMZQrY+EncsrRNjDj7p2HU/GV31xMu9qVZXT7uiLU+srq1CIi4aC0uY0Vp/UNgWMpg8bkZwOUS2qkcjwPc3cCdlukWVNdVE3pm0mpFvFXoz4J4T+HxfZTxV6M+CeE/h8X2VKUW7tGNxzzlbzvRbxV6M+CeE/h8X2U8VejPgnhP4fF9lSlE7Rjcc85LzvRbxV6M+CeE/h8X2U8VejPgnhP4fF9lSlE7Rjcc85LzvaPFaG05grLbOOwGMoWG78s1apHG9u/fsQNxut4iLVVXVXN6pumsREWAIiICIiAiIgIiICIiAiIgIiICIiD/2Q==",
      "text/plain": "<IPython.core.display.Image object>"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import Image\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.prebuilt import ToolNode, create_react_agent\n",
    "import os\n",
    "from langchain_openai import ChatOpenAI\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "model = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\",\n",
    "    temperature=0, streaming=True,\n",
    ")\n",
    "\n",
    "tools = [get_context]\n",
    "# ToolNode will automatically take care of injecting state into tools\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "checkpointer = MemorySaver()\n",
    "graph = create_react_agent(model, tools, state_schema=State, checkpointer=checkpointer)\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:50:55.561195Z",
     "start_time": "2024-11-11T03:50:54.159516Z"
    }
   },
   "id": "26d9c75d31ad644",
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Use it\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1a81c39f09816d01"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001B[1m Human Message \u001B[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_a77b7315a9b94f319c3f1e)\n",
      " Call ID: call_a77b7315a9b94f319c3f1e\n",
      "  Args:\n",
      "    question: what's the latest news about FooBar\n",
      "=================================\u001B[1m Tool Message \u001B[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company just raised 1 Billion dollars!\n",
      "\n",
      "FooBar company was founded in 2019\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "The latest news about FooBar is that the company just raised 1 Billion dollars! Additionally, I found that FooBar was founded in 2019.\n"
     ]
    }
   ],
   "source": [
    "docs = [\n",
    "    \"FooBar company just raised 1 Billion dollars!\",\n",
    "    \"FooBar company was founded in 2019\",\n",
    "]\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}],\n",
    "    \"docs\": docs,\n",
    "}\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "for chunk in graph.stream(inputs, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:00.116233Z",
     "start_time": "2024-11-11T03:50:55.563458Z"
    }
   },
   "id": "18ba9471f0b1611e",
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Pass shared memory (store) to the graph\n",
    "You might also want to give tools access to memory that is shared across multiple conversations or users. We can do it by passing LangGraph Store to the tools using a different annotation -- InjectedStore.\n",
    "\n",
    "Let's modify our example to save the documents in an in-memory store and retrieve them using get_context tool. We'll also make the documents accessible based on a user ID, so that some documents are only visible to certain users. The tool will then use the user_id provided in the config to retrieve a correct set of documents."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c8ecc9d1bb8cb592"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "doc_store = InMemoryStore()\n",
    "\n",
    "namespace = (\"documents\", \"1\")  # user ID\n",
    "doc_store.put(\n",
    "    namespace, \"doc_0\", {\"doc\": \"FooBar company just raised 1 Billion dollars!\"}\n",
    ")\n",
    "namespace = (\"documents\", \"2\")  # user ID\n",
    "doc_store.put(namespace, \"doc_1\", {\"doc\": \"FooBar company was founded in 2019\"})"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:00.125054Z",
     "start_time": "2024-11-11T03:51:00.118493Z"
    }
   },
   "id": "a960dda569e42b70",
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Define the tools"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d51d5a0061500fc8"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "{'description': 'Get relevant context for answering the question.',\n 'properties': {'question': {'title': 'Question', 'type': 'string'}},\n 'required': ['question'],\n 'title': 'get_context',\n 'type': 'object'}"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Tuple\n",
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langgraph.prebuilt import InjectedStore\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_context(\n",
    "    question: str,\n",
    "    config: RunnableConfig,\n",
    "    store: Annotated[BaseStore, InjectedStore()],\n",
    ") -> Tuple[str, List[Document]]:\n",
    "    \"\"\"Get relevant context for answering the question.\"\"\"\n",
    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
    "    docs = [item.value[\"doc\"] for item in store.search((\"documents\", user_id))]\n",
    "    return \"\\n\\n\".join(doc for doc in docs)\n",
    "\n",
    "get_context.tool_call_schema.schema()\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:00.158746Z",
     "start_time": "2024-11-11T03:51:00.127303Z"
    }
   },
   "id": "57f5fb84cc5323b4",
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Define the graph"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "827d19cfdac8c620"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "tools = [get_context]\n",
    "\n",
    "# ToolNode will automatically take care of injecting Store into tools\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "checkpointer = MemorySaver()\n",
    "# NOTE: we need to pass our store to `create_react_agent` to make sure our graph is aware of it\n",
    "graph = create_react_agent(model, tools, checkpointer=checkpointer, store=doc_store)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:00.174204Z",
     "start_time": "2024-11-11T03:51:00.161845Z"
    }
   },
   "id": "7db260712d87bc21",
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Use it"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7dd5a4df8f994dea"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001B[1m Human Message \u001B[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_f63e54733c434bb199b089)\n",
      " Call ID: call_f63e54733c434bb199b089\n",
      "  Args:\n",
      "    question: what's the latest news about FooBar\n",
      "=================================\u001B[1m Tool Message \u001B[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company just raised 1 Billion dollars!\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "The latest news about FooBar is that the company just raised 1 Billion dollars!\n"
     ]
    }
   ],
   "source": [
    "messages = [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}]\n",
    "config = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "for chunk in graph.stream({\"messages\": messages}, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:03.313963Z",
     "start_time": "2024-11-11T03:51:00.182176Z"
    }
   },
   "id": "5ec8623fa96f513c",
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001B[1m Human Message \u001B[0m=================================\n",
      "\n",
      "what's the latest news about FooBar\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "Tool Calls:\n",
      "  get_context (call_7442873abe294318a9a7c8)\n",
      " Call ID: call_7442873abe294318a9a7c8\n",
      "  Args:\n",
      "    question: what's the latest news about FooBar\n",
      "=================================\u001B[1m Tool Message \u001B[0m=================================\n",
      "Name: get_context\n",
      "\n",
      "FooBar company was founded in 2019\n",
      "==================================\u001B[1m Ai Message \u001B[0m==================================\n",
      "\n",
      "The information I found is that the FooBar company was founded in 2019. It seems there may not have been any recent news updates about it or the source I checked did not have the latest updates. Could you please specify which aspect of FooBar you are interested in, so I can look up more relevant information for you?\n"
     ]
    }
   ],
   "source": [
    "messages = [{\"type\": \"user\", \"content\": \"what's the latest news about FooBar\"}]\n",
    "config = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"2\"}}\n",
    "for chunk in graph.stream({\"messages\": messages}, config, stream_mode=\"values\"):\n",
    "    chunk[\"messages\"][-1].pretty_print()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:51:10.763521Z",
     "start_time": "2024-11-11T03:51:03.315329Z"
    }
   },
   "id": "9042384614fbc1d3",
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "StateSnapshot(values={'messages': [HumanMessage(content=\"what's the latest news about FooBar\", additional_kwargs={}, response_metadata={}, id='499626bf-1c2f-4894-8587-79de87a07fa0'), AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_7442873abe294318a9a7c8', 'function': {'arguments': '{\"question\": \"what\\'s the latest news about FooBar\"}', 'name': 'get_context'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls', 'model_name': 'qwen-max'}, id='run-34863dcf-0197-46d0-8431-49926fb64e6e-0', tool_calls=[{'name': 'get_context', 'args': {'question': \"what's the latest news about FooBar\"}, 'id': 'call_7442873abe294318a9a7c8', 'type': 'tool_call'}]), ToolMessage(content='FooBar company was founded in 2019', name='get_context', id='c4d521cc-9e56-45c8-afc9-96dedcad150d', tool_call_id='call_7442873abe294318a9a7c8'), AIMessage(content='The information I found is that the FooBar company was founded in 2019. It seems there may not have been any recent news updates about it or the source I checked did not have the latest updates. Could you please specify which aspect of FooBar you are interested in, so I can look up more relevant information for you?', additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'qwen-max'}, id='run-96d667fa-ac97-479b-bf01-7676d16f15ca-0')]}, next=(), config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef9fe03-0898-6a18-8003-d7ea72628e9e'}}, metadata={'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The information I found is that the FooBar company was founded in 2019. It seems there may not have been any recent news updates about it or the source I checked did not have the latest updates. Could you please specify which aspect of FooBar you are interested in, so I can look up more relevant information for you?', additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'qwen-max'}, id='run-96d667fa-ac97-479b-bf01-7676d16f15ca-0')]}}, 'step': 3, 'parents': {}}, created_at='2024-11-11T03:51:10.758513+00:00', parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1ef9fe02-d0af-6340-8002-be0651ddd413'}}, tasks=())"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_state(config)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-11T03:52:23.784914Z",
     "start_time": "2024-11-11T03:52:23.773512Z"
    }
   },
   "id": "3ce336c26b6fddc5",
   "execution_count": 15
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
